# docx_to_json_fixed_encoding.py
import json
import argparse
import sys
from pathlib import Path
from docx import Document
from docx.shared import RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
import traceback

class DocxToJsonConverter:
    def docx_to_json(self, docx_path: str, output_json: str = None) -> dict:
        """将 DOCX 转换为 JSON - 修复编码版本"""
        try:
            print(f"正在读取文件: {docx_path}")

            # 确保使用正确的编码打开文件
            doc = Document(docx_path)
            print("文件读取成功，开始提取内容...")

            document_data = {
                "metadata": self._extract_metadata(doc),
                "content": self._extract_content(doc),
                "tables": self._extract_tables(doc),
            }

            print("内容提取完成")

            if output_json:
                # 确保使用 UTF-8 编码保存 JSON
                with open(output_json, 'w', encoding='utf-8') as f:
                    json.dump(document_data, f, ensure_ascii=False, indent=2)
                print(f"JSON 文件已保存: {output_json}")

            return {
                "success": True,
                "data": document_data,
                "error": None
            }

        except Exception as e:
            error_msg = f"处理文档时出错: {str(e)}"
            print(f"错误: {error_msg}")
            traceback.print_exc()
            return {
                "success": False,
                "data": None,
                "error": error_msg
            }

    def _extract_metadata(self, doc: Document) -> dict:
        """提取元数据 - 修复编码"""
        try:
            return {
                "title": self._safe_encode(doc.core_properties.title),
                "author": self._safe_encode(doc.core_properties.author),
                "created": str(doc.core_properties.created) if doc.core_properties.created else "",
                "modified": str(doc.core_properties.modified) if doc.core_properties.modified else "",
            }
        except Exception as e:
            print(f"提取元数据时出错: {e}")
            return {}

    def _safe_encode(self, text):
        """安全编码文本"""
        if text is None:
            return ""
        try:
            # 确保文本是字符串并正确处理中文
            return str(text)
        except:
            return ""

    def _extract_content(self, doc: Document) -> list:
        """提取内容 - 修复编码"""
        content = []

        for i, paragraph in enumerate(doc.paragraphs):
            try:
                if paragraph.text and paragraph.text.strip():
                    # 确保中文正确编码
                    para_text = self._safe_encode(paragraph.text)

                    para_data = {
                        "index": i,
                        "text": para_text,
                        "style": self._safe_encode(paragraph.style.name) if paragraph.style and paragraph.style.name else "Normal",
                        "alignment": self._get_alignment(paragraph.alignment),
                        "runs": self._extract_runs(paragraph)
                    }
                    content.append(para_data)
            except Exception as e:
                print(f"提取段落 {i} 时出错: {e}")
                continue

        return content

    def _extract_runs(self, paragraph) -> list:
        """提取运行格式 - 修复编码"""
        runs = []

        for run in paragraph.runs:
            try:
                if run.text and run.text.strip():
                    # 确保中文正确编码
                    run_text = self._safe_encode(run.text)

                    run_data = {
                        "text": run_text,
                        "bold": run.bold if run.bold is not None else False,
                        "italic": run.italic if run.italic is not None else False,
                        "underline": run.underline if run.underline is not None else False,
                    }

                    runs.append(run_data)

            except Exception as e:
                print(f"提取运行时出错: {e}")
                continue

        return runs

    def _extract_tables(self, doc: Document) -> list:
        """提取表格 - 修复编码"""
        tables = []

        for table_idx, table in enumerate(doc.tables):
            try:
                table_data = {
                    "index": table_idx,
                    "rows": []
                }

                for row_idx, row in enumerate(table.rows):
                    row_data = []
                    for cell_idx, cell in enumerate(row.cells):
                        try:
                            # 确保中文正确编码
                            cell_text = self._safe_encode(cell.text)
                            row_data.append({
                                "row": row_idx,
                                "column": cell_idx,
                                "text": cell_text.strip() if cell_text else ""
                            })
                        except Exception as cell_error:
                            print(f"处理单元格时出错: {cell_error}")
                            row_data.append({
                                "row": row_idx,
                                "column": cell_idx,
                                "text": ""
                            })

                    table_data["rows"].append(row_data)

                tables.append(table_data)

            except Exception as e:
                print(f"提取表格 {table_idx} 时出错: {e}")
                continue

        return tables

    def _get_alignment(self, alignment) -> str:
        """获取对齐方式"""
        alignment_map = {
            WD_ALIGN_PARAGRAPH.LEFT: "left",
            WD_ALIGN_PARAGRAPH.CENTER: "center",
            WD_ALIGN_PARAGRAPH.RIGHT: "right",
            WD_ALIGN_PARAGRAPH.JUSTIFY: "justify"
        }
        return alignment_map.get(alignment, "left")

def main():
    # 设置系统编码（Windows 需要）
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

    parser = argparse.ArgumentParser(description='DOCX 转 JSON 转换工具 - 修复编码版本')
    parser.add_argument('--file', required=True, help='输入的 DOCX 文件路径')
    parser.add_argument('--output', '-o', help='输出的 JSON 文件路径')

    args = parser.parse_args()

    converter = DocxToJsonConverter()

    try:
        input_path = Path(args.file)
        if not input_path.exists():
            result = {
                "success": False,
                "data": None,
                "error": f"文件不存在: {args.file}"
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        output_path = args.output
        if not output_path:
            output_path = input_path.with_suffix('.json')

        # 执行转换
        result = converter.docx_to_json(str(input_path), str(output_path))

        # 输出结果
        print(json.dumps(result, ensure_ascii=False))

    except Exception as e:
        error_result = {
            "success": False,
            "data": None,
            "error": f"程序执行错误: {str(e)}"
        }
        print(json.dumps(error_result, ensure_ascii=False))
        sys.exit(1)

if __name__ == "__main__":
    main()